import scala.xml._
import scala.io.Source
import scala.math._
import java.io._

import ACID._

import scala.collection.mutable

object CSV{
  def save(data: Iterable[Iterable[String]], path: String): Unit = {
    val writer = new PrintWriter(new File(path))
    data.foreach(entry => writer.write(entry.reduceLeft(_+", "+_) + "\n"))
    writer.close()
  }
}

object ACID{
  val acid: Seq[String] = Seq("TYR", "ILE", "GLY", "HIS", "MET", "LYS", "ARG", "ASP", "ALA", "GLU", "ASN", "PHE", "THR", "SER", "PRO", "VAL", "LEU", "TRP", "CYS", "GLN")
}

object Analysis {
  private def getCartn(x: Seq[String]): Seq[String] = x match {
    case _ :: _ :: _ :: _ :: _ :: _ :: a => a
    case _ => Seq("0", "0", "0")
  }
  private def distance(x: Seq[String], y: Seq[String]): Double = sqrt((0.0 /: x.zip(y)) ((sum, tup) => sum + pow(tup._1.toDouble - tup._2.toDouble, 2)))

  def analysis(path: String, threshold: Double): Seq[(Seq[Seq[String]], Seq[Seq[String]])] = {
    val source = Source.fromFile(path)
    val xml = XML.load(source.reader())
    val site = (xml \ "atom_siteCategory" \ "atom_site").map(entry => Seq((entry \\ "@id").text, (entry \ "type_symbol").text, (entry \ "auth_atom_id").text, (entry \ "label_comp_id").text, (entry \ "label_entity_id").text, (entry \ "label_seq_id").text, (entry \ "Cartn_x").text, (entry \ "Cartn_y").text, (entry \ "Cartn_z").text))
    val peptide = (xml \ "entity_poly_seqCategory" \ "entity_poly_seq").map(entry => Seq((entry \\ "@entity_id").text, (entry \\ "@mon_id").text, (entry \\ "@num").text))
    val cas = site.filter(_ (2) == "CA")
    val a = Seq.tabulate(cas.map(_ (4).toInt).max)(index => cas.filter(_ (4).toInt == index + 1))
    val (sh3, pep) = a.partition(_.size > a.map(_.size).min)
    sh3.map(_.partition(ele => pep.map(_.map(pepele => distance(getCartn(ele), getCartn(pepele)) < threshold).reduceLeft(_ || _)).reduceLeft(_ || _)))
  }

  def statistics(a: Iterable[(Iterable[Seq[String]], Iterable[Seq[String]])]): Seq[Seq[String]] ={
    val activeMap = a.map(_._1).flatMap(_.map(_(3))).map((_, 1)).groupBy(_._1).map(tup => (tup._1, tup._2.size))
    val deactiveMap = a.map(_._2).flatMap(_.map(_(3))).map((_, 1)).groupBy(_._1).map(tup => (tup._1, tup._2.size))
    val asum = activeMap.values.sum.toDouble
    val scaleactive = activeMap.map(tup => (tup._1, tup._2 / asum))
    val dsum = deactiveMap.values.sum.toDouble
    val scaledeactive = deactiveMap.map(tup => (tup._1, tup._2 / dsum))
    var mapsum:mutable.Map[String, Double] = mutable.Map()
    acid.foreach(f => mapsum += (f -> (activeMap.getOrElse(f, 0) + deactiveMap.getOrElse(f, 0)) / (asum + dsum)))
    acid.map(name => Seq(name, scaleactive.getOrElse(name, 0).toString, scaledeactive.getOrElse(name, 0).toString, mapsum.getOrElse(name, 0).toString))
  }

  def allstatistics(a: Iterable[Iterable[(Iterable[Seq[String]], Iterable[Seq[String]])]]): Seq[Seq[String]] ={
    val activeMap = a.map(_.map(_._1)).flatMap(_.flatMap(_.map(_(3)))).map((_, 1)).groupBy(_._1).map(tup => (tup._1, tup._2.size))
    val deactiveMap = a.map(_.map(_._2)).flatMap(_.flatMap(_.map(_ (3)))).map((_, 1)).groupBy(_._1).map(tup => (tup._1, tup._2.size))
    val asum = activeMap.values.sum.toDouble
    val scaleactive = activeMap.map(tup => (tup._1, tup._2 / asum))
    val dsum = deactiveMap.values.sum.toDouble
    val scaledeactive = deactiveMap.map(tup => (tup._1, tup._2 / dsum))
    var mapsum:mutable.Map[String, Double] = mutable.Map()
    acid.foreach(f => mapsum += (f -> (activeMap.getOrElse(f, 0) + deactiveMap.getOrElse(f, 0)) / (asum + dsum)))
    acid.map(name => Seq(name, scaleactive.getOrElse(name, 0).toString, scaledeactive.getOrElse(name, 0).toString, mapsum.getOrElse(name, 0).toString))
  }

  def main(args: Array[String]): Unit = {
    val path: File = new File("./PDB")
    val name = path.listFiles.map(_.toString).filter(_.endsWith("xml"))
    val a = name.map(path => {
      val temp = analysis(path, 8.0)
      println(path+" done")
      temp
    })
    val b =a.map(statistics(_).map(entry => (entry.head, entry.tail)).toMap)
    CSV.save((name.zip(b.map(entry => acid.map(entry.getOrElse(_, Seq("")).head))).map(tup => tup._2.toList.::(tup._1)).toList.::(acid.toList.::("name"))), "./Ansa.csv")
    CSV.save((name.zip(b.map(entry => acid.map(entry.getOrElse(_, Seq("", "")).tail.head))).map(tup => tup._2.toList.::(tup._1)).toList.::(acid.toList.::("name"))), "./Ansb.csv")
    CSV.save((name.zip(b.map(entry => acid.map(entry.getOrElse(_, Seq("", "", "")).tail.tail.head))).map(tup => tup._2.toList.::(tup._1)).toList.::(acid.toList.::("name"))), "./Ansc.csv")
    CSV.save(allstatistics(a), "./Ans.csv")
  }
}
